import pandas as pd
import os
from datetime import timedelta


def convert_daily_to_weekly(input_path, output_path):
    """将日线数据转换为周线数据"""
    try:
        # 尝试多种编码方式读取
        df = None
        encodings = ['gbk', 'utf-8-sig', 'latin1']

        for encoding in encodings:
            try:
                df = pd.read_csv(input_path, encoding=encoding)
                print(f"成功使用 {encoding} 编码读取文件 {input_path}")
                break
            except UnicodeDecodeError:
                continue

        # 如果前面的编码尝试都失败
        if df is None:
            df = pd.read_csv(input_path, encoding='gbk', errors='replace')
            print(f"使用替换模式读取文件 {input_path}")

    except Exception as e:
        print(f"读取文件 {input_path} 时发生错误: {str(e)}")
        return False

    # 验证必要字段
    required_columns = ['股票代码', '交易日期', '后复权开盘价', '后复权收盘价',
                        '后复权最高价', '后复权最低价', '成交量', '成交额']

    if not all(col in df.columns for col in required_columns):
        print(f"文件 {input_path} 缺少必要字段，跳过处理")
        return False

    # 强制转换日期格式
    def parse_date(date_str):
        try:
            date_str = str(date_str).strip()
            if date_str.isdigit() and len(date_str) == 8:
                return pd.to_datetime(date_str, format='%Y%m%d')
            return pd.to_datetime(date_str)
        except:
            return pd.NaT

    df['交易日期'] = df['交易日期'].apply(parse_date)

    # 排除无效日期
    df = df[df['交易日期'].notna()].copy()
    if df.empty:
        print(f"文件 {input_path} 中没有有效日期数据")
        return False

    # 按日期排序
    df.sort_values('交易日期', inplace=True)
    df.reset_index(drop=True, inplace=True)

    # 处理数据
    weekly_data = []
    current_group = []

    for _, row in df.iterrows():
        if not current_group:
            current_group.append(row)
        else:
            prev_date = current_group[-1]['交易日期']
            curr_date = row['交易日期']

            if (curr_date - prev_date).days == 1:
                current_group.append(row)
            else:
                if prev_date.weekday() == 4:  # 周五
                    weekly_data.append(process_group(current_group))
                    current_group = [row]
                else:
                    prev_week = prev_date.isocalendar().week
                    curr_week = curr_date.isocalendar().week

                    if prev_week != curr_week:
                        weekly_data.append(process_group(current_group))
                        current_group = [row]
                    else:
                        current_group.append(row)

    # 处理最后一组数据
    if current_group:
        weekly_data.append(process_group(current_group))

    # 创建DataFrame
    if not weekly_data:
        print(f"文件 {input_path} 未生成任何周线数据")
        return False

    weekly_df = pd.DataFrame(weekly_data)
    columns = ['股票代码', '交易开始日期', '交易结束日期',
               '周线_后复权开盘价', '周线_后复权最高价',
               '周线_后复权最低价', '周线_后复权收盘价',
               '周线_成交量', '周线_成交额']

    weekly_df = weekly_df[columns].copy()

    # 格式化数值列（保留两位小数）
    numeric_cols = ['周线_后复权开盘价', '周线_后复权最高价',
                    '周线_后复权最低价', '周线_后复权收盘价',
                    '周线_成交量', '周线_成交额']

    weekly_df[numeric_cols] = weekly_df[numeric_cols].round(2)

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # 保存结果
    try:
        weekly_df.to_csv(output_path, index=False, encoding='utf-8-sig')
        print(f"成功处理: {input_path}")
        return True
    except Exception as e:
        print(f"保存文件 {output_path} 时发生错误: {str(e)}")
        return False


def process_group(group):
    """处理一组连续交易日的数据"""
    first_row = group[0]
    last_row = group[-1]

    return {
        '股票代码': first_row['股票代码'],
        '交易开始日期': first_row['交易日期'],
        '交易结束日期': last_row['交易日期'],
        '周线_后复权开盘价': first_row['后复权开盘价'],
        '周线_后复权收盘价': last_row['后复权收盘价'],
        '周线_后复权最高价': max(row['后复权最高价'] for row in group),
        '周线_后复权最低价': min(row['后复权最低价'] for row in group),
        '周线_成交量': sum(row['成交量'] for row in group),
        '周线_成交额': sum(row['成交额'] for row in group)
    }


def batch_convert_daily_to_weekly(input_dir, output_dir):
    """批量处理目录下的所有CSV文件"""
    success_count = 0
    fail_count = 0

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 遍历目录下的所有CSV文件
    for filename in os.listdir(input_dir):
        if filename.lower().endswith('.csv'):
            input_path = os.path.join(input_dir, filename)
            output_filename = f"周线{filename}"
            output_path = os.path.join(output_dir, output_filename)

            try:
                result = convert_daily_to_weekly(input_path, output_path)
                if result:
                    success_count += 1
                else:
                    fail_count += 1
            except Exception as e:
                print(f"处理文件 {filename} 时发生异常: {str(e)}")
                fail_count += 1

    print(f"\n处理完成！成功处理: {success_count} 个文件，失败: {fail_count} 个文件")


# 使用示例
if __name__ == "__main__":
    input_directory = r"D:\pythonProject\BaiduSyncdisk\LongMoneyTrade\各部分数据\后复权日线\个股_后复权日线"
    output_directory = r"D:\pythonProject\BaiduSyncdisk\LongMoneyTrade\各部分数据\后复权周线"

    batch_convert_daily_to_weekly(input_directory, output_directory)
